package com.xugui.learn.alibaba.hook;

import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.alibaba.cloud.ai.graph.agent.hook.HookPosition;
import com.alibaba.cloud.ai.graph.agent.hook.ModelHook;
import org.springframework.ai.chat.messages.Message;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;

/**
 * ModelHook - 在模型调用前后执行（例如：消息修剪），区别于 AgentHook
 * ModelHook在一次 agent 调用中可能会调用多次，也就是每次 reasoning-acting 迭代都会执行
 */
// @HookPositions({HookPosition.BEFORE_MODEL, HookPosition.AFTER_MODEL})
public class MessageTrimmingHook extends ModelHook {
    private static final int MAX_MESSAGES = 10;

    @Override
    public String getName() {
        return "message_trimming";
    }

    @Override
    public HookPosition[] getHookPositions() {
        return new HookPosition[]{HookPosition.BEFORE_MODEL, HookPosition.AFTER_MODEL};
    }

    @Override
    public CompletableFuture<Map<String, Object>> beforeModel(OverAllState state, RunnableConfig config) {
        Optional<Object> messagesOpt = state.value("messages");
        if (messagesOpt.isPresent()) {
            @SuppressWarnings("unchecked")
            List<Message> messages = (List<Message>) messagesOpt.get();
            if (messages.size() > MAX_MESSAGES) {
                return CompletableFuture.completedFuture(Map.of("messages",
                        messages.subList(messages.size() - MAX_MESSAGES, messages.size())));
            }
        }
        return CompletableFuture.completedFuture(Map.of());
    }

    @Override
    public CompletableFuture<Map<String, Object>> afterModel(OverAllState state, RunnableConfig config) {
        return CompletableFuture.completedFuture(Map.of());
    }
}